import streamlit as st
from chatbot.chat_core import generate_response

# 设置页面标题
st.title("多服务聊天机器人")

# 选择服务
service_options = ["silicon", "huawei", "deepseek"]
selected_service = st.selectbox("选择服务", service_options)

# 初始化聊天历史
if "messages" not in st.session_state:
    st.session_state.messages = []

# 显示聊天历史
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 获取用户输入
if prompt := st.chat_input("请输入你的问题"):
    # 添加用户消息到聊天历史
    st.session_state.messages.append({"role": "user", "content": prompt})
    # 显示用户消息
    with st.chat_message("user"):
        st.markdown(prompt)

    # 生成 AI 回复
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        for response_chunk in generate_response(st.session_state.messages, selected_service):
            full_response += response_chunk
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)

    # 添加 AI 回复到聊天历史
    st.session_state.messages.append({"role": "assistant", "content": full_response})